#include "Trie.h"

#include <stack>

template class DS_API nmspc::Shield<std::vector<std::string>>;

nmspc::Trie::Trie() :children(std::unordered_map<char, std::unique_ptr<Trie>>()), if_end(false)
{
}

nmspc::Trie::~Trie()
{
}

nmspc::Trie* nmspc::Trie::search(const std::string& prefix)
{
	nmspc::Trie* node = this;
	for (const char& ch : prefix)
	{
		if (node->children.find(ch) == node->children.end())
		{
			return nullptr;
		}
		node = node->children[ch].get();
	}
	return node;
}

bool nmspc::Trie::inserted(const std::string& whole_word)
{
	nmspc::Trie* node = this->search(whole_word);
	return node && node->if_end;
}

void nmspc::Trie::insert(const std::string& word)
{
	nmspc::Trie* node = this;
	for (const char& ch : word)
	{
		if (node->children.find(ch) == node->children.end())
		{
			node->children.emplace(ch, std::move(std::make_unique<Trie>()));
		}
		node = node->children[ch].get();
	}
	node->if_end = true;
}

void nmspc::Trie::erase(const std::string& word)
{
	nmspc::Trie* node = this;
	std::stack<nmspc::Trie*> nodes;
	std::stack<char> edges;
	for (const char& ch : word)
	{
		if (node->children.find(ch) == node->children.end())
		{
			return;
		}
		nodes.push(node);
		edges.push(ch);
		node = node->children[ch].get();
	}
	if (nodes.empty() || !nodes.top()->children[edges.top()]->if_end)
	{
		return;
	}
	else
	{
		nodes.top()->children[edges.top()]->if_end = false;
	}
	while (!nodes.empty())
	{
		Trie* parent = nodes.top();
		char edge = edges.top();
		nodes.pop();
		edges.pop();
		if (parent->children[edge]->children.empty() && !parent->children[edge]->if_end)
		{
			parent->children.erase(edge);
		}
		else
		{
			break;
		}
	}
}

void nmspc::Trie::getSuffix(nmspc::SVectorString& all_suffix, std::string& cur_suffix)
{
	if (this->if_end == true)
	{
		if (cur_suffix != "") all_suffix.get().push_back(cur_suffix);
	}
	for (const auto& [ch, child] : this->children)
	{
		cur_suffix += ch;
		child->getSuffix(all_suffix, cur_suffix);
		cur_suffix.pop_back();
	}
}

nmspc::SVectorString nmspc::Trie::getSuffix(const std::string& prefix)
{
	nmspc::Trie* node = this->search(prefix);
	if (!node) return {};
	std::vector<std::string> all_suffix = std::vector<std::string>();
	std::stack<std::pair<nmspc::Trie*, std::string>> dfs_stack;
	dfs_stack.push({ node, "" });
	while (!dfs_stack.empty()) 
	{
		auto [node, cur_suffix] = dfs_stack.top();
		dfs_stack.pop();
		if (node->if_end && !cur_suffix.empty()) 
		{
			all_suffix.push_back(cur_suffix);
		}
		for (const auto& [ch, child] : node->children) 
		{
			dfs_stack.push({ child.get(), cur_suffix + ch });
		}
	}
	return SVectorString(all_suffix);
}
